import torch
import torch.nn as nn
from torchvision.models import alexnet

# 创建保存目录
import os

os.makedirs('../static', exist_ok=True)

# 加载预训练的AlexNet模型（默认输出1000类）
pretrained_model = alexnet(pretrained=True)

# 打印原始模型结构（可选）
print("原始AlexNet模型结构：")
print(pretrained_model)

# 修改最后一层全连接层（classifier[6]）为45类输出
pretrained_model.classifier[6] = nn.Linear(
    in_features=4096,  # 保持输入维度不变
    out_features=45,  # 修改输出维度为45类
    bias=True
)

# 打印修改后的模型结构（可选）
print("\n修改后的模型结构：")
print(pretrained_model)

# 保存修改后的模型参数
save_path = "../static/AlexNet_mammals_45classes.pth"
torch.save(pretrained_model.state_dict(), save_path)
print(f"\n已保存适配45类的模型到：{save_path}")

# 验证保存的模型是否可以正确加载
try:
    # 创建一个新的45类AlexNet模型
    model = alexnet(num_classes=45)

    # 加载保存的参数
    loaded_state_dict = torch.load(save_path)
    model.load_state_dict(loaded_state_dict)

    print("模型加载验证成功！")
except Exception as e:
    print(f"模型加载验证失败: {e}")